package com.iot.protocol;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class UdpServer implements AutoCloseable {
    private final int port;
    private final DatagramSocket serverSocket;
    private final ExecutorService executorService;
    private final AtomicInteger activeConnections;
    private volatile boolean running;
    private static final int BUFFER_SIZE = 1024;

    public UdpServer(int port) throws IOException {
        this.port = port;
        this.serverSocket = new DatagramSocket(port);
        this.executorService = Executors.newCachedThreadPool();
        this.activeConnections = new AtomicInteger(0);
        this.running = true;
        log.info("UDP服务器已启动，监听端口: {}", port);
    }

    public void start() {
        try {
            byte[] receiveBuffer = new byte[BUFFER_SIZE];
            DatagramPacket receivePacket = new DatagramPacket(receiveBuffer, receiveBuffer.length);

            while (running) {
                try {
                    // 接收数据包
                    serverSocket.receive(receivePacket);
                    activeConnections.incrementAndGet();

                    // 获取客户端信息
                    InetAddress clientAddress = receivePacket.getAddress();
                    int clientPort = receivePacket.getPort();
                    String clientInfo = String.format("%s:%d", clientAddress.getHostAddress(), clientPort);

                    // 处理接收到的数据
                    String receivedData = new String(receivePacket.getData(), 0, receivePacket.getLength(), StandardCharsets.UTF_8);
                    log.info("收到来自 {} 的消息: {}", clientInfo, receivedData);

                    // 异步处理响应
                    executorService.execute(() -> handleClientMessage(clientAddress, clientPort, receivedData));

                } catch (IOException e) {
                    if (running) {
                        log.error("接收数据包时发生错误: {}", e.getMessage(), e);
                    }
                } finally {
                    activeConnections.decrementAndGet();
                }
            }
        } finally {
            shutdown();
        }
    }

    private void handleClientMessage(InetAddress clientAddress, int clientPort, String message) {
        try {
            // 构造响应消息
            String response = String.format("服务器已收到消息: %s", message);
            byte[] responseData = response.getBytes(StandardCharsets.UTF_8);
            DatagramPacket sendPacket = new DatagramPacket(
                responseData, 
                responseData.length, 
                clientAddress, 
                clientPort
            );

            // 发送响应
            serverSocket.send(sendPacket);
            log.info("已向 {}:{} 发送响应", clientAddress.getHostAddress(), clientPort);
        } catch (IOException e) {
            log.error("发送响应时发生错误: {}", e.getMessage(), e);
        }
    }

    public void broadcast(String message) {
        if (!running) return;

        log.info("广播消息: {}", message);
        byte[] messageData = message.getBytes(StandardCharsets.UTF_8);
        
        try {
            // 广播到所有可能的客户端（这里使用255.255.255.255作为广播地址）
            DatagramPacket broadcastPacket = new DatagramPacket(
                messageData,
                messageData.length,
                InetAddress.getByName("255.255.255.255"),
                port
            );
            serverSocket.send(broadcastPacket);
            log.info("广播消息已发送");
        } catch (IOException e) {
            log.error("广播消息时发生错误: {}", e.getMessage(), e);
        }
    }

    public void showStatus() {
        System.out.printf("当前状态：\n" +
                "监听端口：%d\n" +
                "活跃连接数：%d\n" +
                "服务器状态：%s\n",
                port,
                activeConnections.get(),
                running ? "运行中" : "已停止");
    }

    @Override
    public void close() {
        shutdown();
    }

    private void shutdown() {
        running = false;
        log.info("正在关闭UDP服务器...");

        // 关闭线程池
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }

        // 关闭服务器socket
        if (serverSocket != null && !serverSocket.isClosed()) {
            serverSocket.close();
        }

        log.info("UDP服务器已关闭");
    }

    public static void main(String[] args) {
        int port = 57673; // 默认端口

        if (args.length >= 1) {
            port = Integer.parseInt(args[0]);
        }

        try (UdpServer server = new UdpServer(port)) {
            // 启动服务器
            server.start();
        } catch (IOException e) {
            log.error("启动服务器失败: {}", e.getMessage(), e);
        }
    }
} 